package com.google.code.securityflow.util

object MultiMap {
  def apply[K, V](key: K, value: V) = new MultiMap[K, V](Map((key, Set(value))));

  def empty[K, V] = new MultiMap[K, V](Map.empty);
}

class MultiMap[K, V](val map: Map[K, Set[V]]) extends (K => Set[V]) with Traversable[(K, V)] {

  def keySet = map.keySet

  override def isEmpty = map.isEmpty

  def valueSet = map.foldLeft[Set[V]](Set.empty)({
    case (b, (k, vs)) => b ++ vs
  })

  def apply(key: K) = get(key);

  def foreach[U](f: ((K, V)) => U) {
    for ((k, vs) <- map; v <- vs) {
      f(k, v);
    }
  }

  def get(key: K): Set[V] = {
    map.get(key) match {
      case None => Set.empty
      case Some(vs) => vs
    }
  }

  def contains(key: K, value: V) = {
    map.get(key) match {
      case None => false
      case Some(vs) => vs.contains(value)
    }
  }

  def add(key: K, value: V): MultiMap[K, V] = {
    map.get(key) match {
      case None => {
        new MultiMap[K, V](map + (key -> Set(value)));
      }
      case Some(vs) => {
        new MultiMap[K, V](map + (key -> (vs + value)));
      }
    }
  }

  def addAll(that: MultiMap[K, V]): MultiMap[K, V] = {
    // TODO: Low Priority: Improve implementation
    new MultiMap[K, V]((for (key <- (this.keySet ++ that.keySet)) yield (key -> (this.map(key) ++ that.map(key)))).toMap)
  }

  def remove(key: K, value: V): MultiMap[K, V] = {
    map.get(key) match {
      case None => this;
      case Some(vs) => {
        if (vs.size == 1 && vs.contains(value)) {
          new MultiMap[K, V](map - key);
        }
        new MultiMap[K, V](map + (key -> (vs - value)));
      }
    }
  }

  def removeKey(key: K): MultiMap[K, V] = {
    new MultiMap[K, V](map - key);
  }

  def removeAll(that: MultiMap[K, V]): MultiMap[K, V] = {
    // TODO: Low Priority: Improve implementation
    val thisKeySet = this.keySet;
    val thatKeySet = that.keySet;

    var newMap: Map[K, Set[V]] = Map.empty;
    for (key <- thatKeySet) {
      val thisSet = this(key);
      val thatSet = that(key);
      val resultSet = thisSet -- thatSet;
      if (!resultSet.isEmpty) {
        newMap = newMap + (key -> resultSet);
      }
    }
    new MultiMap[K, V](newMap ++ this.map.filterKeys(k => !thatKeySet(k)));
  }

  def containsKey(key: K) = map.contains(key)

  def +(kv: (K, V)) = add(kv._1, kv._2);

  def ++(m: MultiMap[K, V]) = addAll(m);

  def ++(kvs: (K, Set[V])) = kvs._2.foldLeft(this)((m, v) => m.add(kvs._1, v));

  def -(kv: (K, V)) = remove(kv._1, kv._2);

  def -(k: K) = removeKey(k);

  def --(m: MultiMap[K, V]) = removeAll(m);

  def inverse: MultiMap[V, K] = {
    var result = MultiMap.empty[V, K];
    for ((key, vs) <- map; v <- vs) {
      result = result + (v -> key);
    }
    result;
  }

  override def equals(o: Any) = {
    o match {
      case that: MultiMap[K, V] => this.map == that.map;
      case _ => false;
    }
  }

  override def hashCode = map.hashCode();

}


